{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93ae9bad-b8cc-43de-ba7d-387e0155674c",
   "metadata": {},
   "source": [
    "# Contextual Retrieval for Multimodal RAG\n",
    "\n",
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_cloud_services/blob/main/examples/parse/multimodal/multimodal_contextual_retrieval_rag.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "In this cookbook we show you how to build a multimodal RAG pipeline with **contextual retrieval**.\n",
    "\n",
    "Contextual retrieval was initially introduced in this Anthropic [blog post](https://www.anthropic.com/news/contextual-retrieval). The high-level intuition is that every chunk is given a concise summary of where that chunk fits in with respect to the overall summary of the document. This allows insertion of high-level concepts/keywords that enable this chunk to be better retrieved for different types of queries.\n",
    "\n",
    "These LLM calls are expensive. Contextual retrieval depends on **prompt caching** in order to be efficient.\n",
    "\n",
    "In this notebook, we use Claude 3.5-Haiku to generate contextual summaries. We cache the document as text tokens, but generate contextual summaries by feeding in the parsed text chunk. \n",
    "\n",
    "We feed both the text and image chunks into the final multimodal RAG pipeline to generate the response.\n",
    "\n",
    "Status:\n",
    "| Last Executed | Version | State      |\n",
    "|---------------|---------|------------|\n",
    "| Aug-20-2025   | 0.6.61  | Maintained |\n",
    "\n",
    "![mm_rag_diagram](./multimodal_contextual_retrieval_rag_img.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54e8d9a7-5036-4d32-818f-00b2e888521f",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "155afa97",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-cloud-services \"llama-index>=0.13.0<0.14.0\" llama-index-embeddings-voyageai llama-index-llms-anthropic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225c5556-a789-4386-a1ee-cce01dbeb6cf",
   "metadata": {},
   "source": [
    "### (Optional) Setup Observability\n",
    "\n",
    "We setup an integration with LlamaTrace (integration with Arize).\n",
    "\n",
    "If you haven't already done so, make sure to create an account here: https://llamatrace.com/login. Then create an API key and put it in the `PHOENIX_API_KEY` variable below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eabee1f-290a-4c85-b362-54f45c8559ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U llama-index-callbacks-arize-phoenix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaeb245c-730b-4c34-ad68-708fdde0e6cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup Arize Phoenix for logging/observability\n",
    "import llama_index.core\n",
    "import os\n",
    "\n",
    "PHOENIX_API_KEY = \"<PHOENIX_API_KEY>\"\n",
    "os.environ[\"OTEL_EXPORTER_OTLP_HEADERS\"] = f\"api_key={PHOENIX_API_KEY}\"\n",
    "llama_index.core.set_global_handler(\n",
    "    \"arize_phoenix\", endpoint=\"https://llamatrace.com/v1/traces\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb362db-b1b1-4eea-be1a-b1f78b0779d7",
   "metadata": {},
   "source": [
    "### Load Data\n",
    "\n",
    "Here we load the [ICONIQ 2024 State of AI Report](https://cdn.prod.website-files.com/65e1d7fb19a3e64b5c36fb38/66eb856e019e59758ef73759_ICONIQ%20Analytics%20%2B%20Insights%20-%20State%20of%20AI%20Sep24.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bce3407-a7d2-47e8-9eaf-ab297a94750c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir data\n",
    "!mkdir data_images_iconiq\n",
    "!wget \"https://cdn.prod.website-files.com/65e1d7fb19a3e64b5c36fb38/66eb856e019e59758ef73759_ICONIQ%20Analytics%20%2B%20Insights%20-%20State%20of%20AI%20Sep24.pdf\" -O data/iconiq_report.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "246ba6b0-51af-42f9-b1b2-8d3e721ef782",
   "metadata": {},
   "source": [
    "### Model Setup\n",
    "\n",
    "Setup models that will be used for downstream orchestration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2957a512-215e-4ba2-ae1a-95e612e0f72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# replace with your Anthropic API key\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = \"sk-...\"\n",
    "# replace with your VoyageAI key\n",
    "os.environ[\"VOYAGE_API_KEY\"] = \"pa-...\"\n",
    "# replace with your LlamaCloud API key\n",
    "os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"llx-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e2071d-bbc2-4707-8ae7-cb4e1fecafd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/loganmarkewich/llama_parse/py/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.anthropic import Anthropic\n",
    "from llama_index.embeddings.voyageai import VoyageEmbedding\n",
    "from llama_index.core import Settings\n",
    "\n",
    "\n",
    "llm = Anthropic(model=\"claude-4-sonnet-20250514\")\n",
    "embed_model = VoyageEmbedding(model_name=\"voyage-3.5\")\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f6416f-f580-4722-aaa9-7f3500408547",
   "metadata": {},
   "source": [
    "## Use LlamaParse to Parse Text and Images\n",
    "\n",
    "In this example, use LlamaParse to parse both the text and images from the document.\n",
    "\n",
    "We parse out the text with LlamaParse premium.\n",
    "\n",
    "**NOTE**: The report has 40 pages, and at ~5c per page, this will cost you $2. Just a heads up!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "570089e5-238a-4dcc-af65-96e7393c2b4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_cloud_services import LlamaParse\n",
    "\n",
    "\n",
    "parser = LlamaParse(\n",
    "    parse_mode=\"parse_page_with_agent\",\n",
    "    model=\"openai-gpt-4-1-mini\",\n",
    "    high_res_ocr=True,\n",
    "    adaptive_long_table=True,\n",
    "    outlined_table_extraction=True,\n",
    "    output_tables_as_HTML=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef82a985-4088-4bb7-9a21-0318e1b9207d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Started parsing the file under job_id 1384d483-16c8-4b20-a3ff-6863eafecbc1\n"
     ]
    }
   ],
   "source": [
    "results = await parser.aparse(\"data/iconiq_report.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5318fb7b-fe6a-4a8a-b82e-4ed7b4512c37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "# A Decision-Making Framework\n",
      "\n",
      "When making decisions around GenAI investments, we believe it will be important to assess organization readiness, put in place a framework and processes for use case evaluation, and proactively mitigate risks\n",
      "\n",
      "----\n",
      "\n",
      "### Accelerate Value  \n",
      "Find synergies between organizational readiness, use cases, and risk mitigation when making GenAI investment decisions\n",
      "\n",
      "----\n",
      "\n",
      "### Use Case Identification & Evaluation  \n",
      "When determining use cases for GenAI, we believe stakeholders will need to assess business value, the fluency vs. accuracy of solutions, and the level of risk associated. Given the risks involved with using GenAI to build new products, many organizations are first starting with use cases for internal productivity.\n",
      "\n",
      "It is also important to implement feedback loops and a system for measuring ROI to evaluate use cases.\n",
      "\n",
      "----\n",
      "\n",
      "### Organizational Readiness  \n",
      "For enterprises adopting GenAI solutions for the first time, we believe it will be important to ensure various components of the organization are ready to support the development and integration needs involved.  \n",
      "Organizational readiness components to assess could include:\n",
      "\n",
      "* Employee readiness and training  \n",
      "* IT / data team expertise  \n",
      "* Security  \n",
      "* Governance structure and policies  \n",
      "* Data ecosystem maturity  \n",
      "\n",
      "----\n",
      "\n",
      "### Risk Mitigation  \n",
      "We believe enterprises will need to account for various risks like data security and privacy concerns, algorithm accuracy / bias, integration complexity, etc. when evaluating GenAI solutions.\n",
      "\n",
      "Organizations can employ various strategies to mitigate some of these risks. For example, it may make sense to invest in fine-tuning or retrieval augmented generation (RAG) techniques to mitigate concerns of model accuracy.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(results.pages[10].md)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d50913fd",
   "metadata": {},
   "source": [
    "We can download the page screenshots directly, and we can use them as context later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "056ba139",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_nodes = await results.aget_image_nodes(\n",
    "    include_object_images=False,\n",
    "    include_screenshot_images=True,\n",
    "    image_download_dir=\"./iconiq_images\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cda70ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_nodes = results.get_markdown_nodes(split_by_page=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd3e098b-0606-4429-b48d-d4fe0140fc0e",
   "metadata": {},
   "source": [
    "## Build Multimodal Index\n",
    "\n",
    "In this section we build the multimodal index over the parsed deck. \n",
    "\n",
    "We do this by creating **text** nodes from the document that contain metadata referencing the original image path.\n",
    "\n",
    "In this example we're indexing the text node for retrieval. The text node has a reference to both the parsed text as well as the image screenshot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "346fe5ef-171e-4a54-9084-7a7805103a13",
   "metadata": {},
   "outputs": [],
   "source": [
    "for text_node, image_node in zip(text_nodes, image_nodes):\n",
    "    text_node.metadata[\"image_path\"] = image_node.image_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32c13950-c1db-435f-b5b4-89d62b8b7744",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "page_number: 1\n",
      "file_name: data/iconiq_report.pdf\n",
      "image_path: iconiq_images/page_1.jpg\n",
      "\n",
      "\n",
      "# The State of AI\n",
      "\n",
      "September 2024\n",
      "\n",
      "Navigating the present and promise of Generative AI\n",
      "\n",
      "ICONIQ | Growth\n"
     ]
    }
   ],
   "source": [
    "print(text_nodes[0].get_content(metadata_mode=\"all\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b57d0fd-e026-45d9-bbe0-c5058be24b6b",
   "metadata": {},
   "source": [
    "#### Add Contextual Summaries\n",
    "\n",
    "In this section we implement the key step in contextual retrieval - attaching metadata to each chunk that situates it within the overall document context.\n",
    "\n",
    "We take advantage of prompt caching by feeding in the static document as prefix tokens, and only swap out the \"header\" tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b96842f9-5c2d-4dd4-9574-c0b56d93ec27",
   "metadata": {},
   "outputs": [],
   "source": [
    "from copy import deepcopy\n",
    "from llama_index.core.llms import ChatMessage, TextBlock, ImageBlock, CachePoint\n",
    "import time\n",
    "\n",
    "\n",
    "whole_doc_text = \"\"\"\\\n",
    "Here is the entire document.\n",
    "<document>\n",
    "{WHOLE_DOCUMENT}\n",
    "</document>\"\"\"\n",
    "\n",
    "chunk_text = \"\"\"\\\n",
    "Here is the chunk we want to situate within the whole document\n",
    "<chunk>\n",
    "{CHUNK_CONTENT}\n",
    "</chunk>\"\"\"\n",
    "\n",
    "suffix_text = \"\"\"Please give a short succinct context to situate this chunk within the overall document for \\\n",
    "the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.\"\"\"\n",
    "\n",
    "\n",
    "def create_contextual_nodes(nodes, llm):\n",
    "    \"\"\"Function to create contextual nodes for a list of nodes\"\"\"\n",
    "    nodes_modified = []\n",
    "\n",
    "    # get overall doc_text string\n",
    "    doc_text = \"\\n\".join([n.get_content(metadata_mode=\"all\") for n in nodes])\n",
    "\n",
    "    for idx, node in enumerate(nodes):\n",
    "        start_time = time.time()\n",
    "        new_node = deepcopy(node)\n",
    "\n",
    "        messages = [\n",
    "            ChatMessage(\n",
    "                role=\"user\",\n",
    "                blocks=[\n",
    "                    TextBlock(text=whole_doc_text.format(WHOLE_DOCUMENT=doc_text)),\n",
    "                    CachePoint(cache_control={\"type\": \"ephemeral\"}),\n",
    "                    TextBlock(\n",
    "                        text=chunk_text.format(\n",
    "                            CHUNK_CONTENT=node.get_content(metadata_mode=\"all\")\n",
    "                        )\n",
    "                    ),\n",
    "                    TextBlock(\n",
    "                        text=\"And here is the page screenshot for the corresponding chunk:\"\n",
    "                    ),\n",
    "                    ImageBlock(path=node.metadata[\"image_path\"]),\n",
    "                    TextBlock(text=suffix_text),\n",
    "                ],\n",
    "            ),\n",
    "        ]\n",
    "\n",
    "        new_response = llm.chat(messages)\n",
    "        new_node.metadata[\"context\"] = str(new_response)\n",
    "\n",
    "        nodes_modified.append(new_node)\n",
    "        print(f\"Completed node {idx}, {time.time() - start_time}\")\n",
    "\n",
    "    return nodes_modified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50c9cfe9-091f-41b7-a7b0-1bec18ac3ea6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed node 0, 5.0501158237457275\n",
      "Completed node 1, 4.125281095504761\n",
      "Completed node 2, 3.700598955154419\n",
      "Completed node 3, 4.249290943145752\n",
      "Completed node 4, 4.552713871002197\n",
      "Completed node 5, 3.700002908706665\n",
      "Completed node 6, 4.9324049949646\n",
      "Completed node 7, 6.246585845947266\n",
      "Completed node 8, 5.678989887237549\n",
      "Completed node 9, 4.55932092666626\n",
      "Completed node 10, 4.865902662277222\n",
      "Completed node 11, 4.376728057861328\n",
      "Completed node 12, 3.823659896850586\n",
      "Completed node 13, 4.069238185882568\n",
      "Completed node 14, 3.7528319358825684\n",
      "Completed node 15, 3.789531946182251\n",
      "Completed node 16, 4.54377818107605\n",
      "Completed node 17, 3.3560800552368164\n",
      "Completed node 18, 4.519093990325928\n",
      "Completed node 19, 5.594789028167725\n",
      "Completed node 20, 3.7624330520629883\n",
      "Completed node 21, 3.778661012649536\n",
      "Completed node 22, 3.895768880844116\n",
      "Completed node 23, 3.6451258659362793\n",
      "Completed node 24, 9.422847032546997\n",
      "Completed node 25, 3.954685926437378\n",
      "Completed node 26, 3.4985830783843994\n",
      "Completed node 27, 3.368708848953247\n",
      "Completed node 28, 3.9136807918548584\n",
      "Completed node 29, 3.791595935821533\n",
      "Completed node 30, 3.1155011653900146\n",
      "Completed node 31, 3.9999842643737793\n",
      "Completed node 32, 3.654320001602173\n",
      "Completed node 33, 3.854135036468506\n",
      "Completed node 34, 3.843966007232666\n",
      "Completed node 35, 4.019424915313721\n",
      "Completed node 36, 9.035747766494751\n",
      "Completed node 37, 5.066689968109131\n",
      "Completed node 38, 7.529208660125732\n",
      "Completed node 39, 4.811733961105347\n",
      "Completed node 40, 2.8257930278778076\n"
     ]
    }
   ],
   "source": [
    "context_llm = Anthropic(model=\"claude-3-5-haiku-latest\")\n",
    "\n",
    "new_text_nodes = create_contextual_nodes(text_nodes, context_llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f404f56-db1e-4ed7-9ba1-ead763546348",
   "metadata": {},
   "source": [
    "#### Build Index\n",
    "\n",
    "Once the text nodes are ready, we feed into our vector store index abstraction, which will index these nodes into a simple in-memory vector store (of course, you should definitely check out our 40+ vector store integrations!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ea53c31-0e38-421c-8d9b-0e3adaa1677e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex(nodes=new_text_nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ea66d70-df0d-42a2-adb3-b20c9d29878a",
   "metadata": {},
   "source": [
    "#### Build Baseline Index\n",
    "\n",
    "Build a baseline index with the text nodes without summarized context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b1f0efa-77ca-4e5a-9d48-17dc442615f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_index = VectorStoreIndex(text_nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f0e33a4-9422-498d-87ee-d917bdf74d80",
   "metadata": {},
   "source": [
    "## Build Multimodal Query Engine\n",
    "\n",
    "We now use LlamaIndex abstractions to build a **custom query engine**. In contrast to a standard RAG query engine that will retrieve the text node and only put that into the prompt (response synthesis module), this custom query engine will also load the image document, and put both the text and image document into the response synthesis module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35a94be2-e289-41a6-92e4-d3cb428fb0c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.query_engine import CustomQueryEngine\n",
    "from llama_index.core.retrievers import BaseRetriever\n",
    "from llama_index.core.schema import MetadataMode\n",
    "from llama_index.core.base.response.schema import Response\n",
    "\n",
    "\n",
    "qa_prompt_block_text = \"\"\"\\\n",
    "Below we give parsed text from slides in two different formats, as well as the image.\n",
    "\n",
    "---------------------\n",
    "{context_str}\n",
    "---------------------\n",
    "\"\"\"\n",
    "\n",
    "image_prefix_block = TextBlock(text=\"And here are the corresponding images per page\\n\")\n",
    "\n",
    "image_suffix = \"\"\"\\\n",
    "Given the context information and not prior knowledge, answer the query. Explain whether you got the answer\n",
    "from the parsed markdown or raw text or image, and if there's discrepancies, and your reasoning for the final answer.\n",
    "\n",
    "Query: {query_str}\n",
    "Answer: \"\"\"\n",
    "\n",
    "\n",
    "class MultimodalQueryEngine(CustomQueryEngine):\n",
    "    \"\"\"Custom multimodal Query Engine.\n",
    "\n",
    "    Takes in a retriever to retrieve a set of document nodes and respond using an LLM + retrieved text/images.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    retriever: BaseRetriever\n",
    "    llm: Anthropic\n",
    "\n",
    "    def __init__(self, **kwargs) -> None:\n",
    "        \"\"\"Initialize.\"\"\"\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def custom_query(self, query_str: str):\n",
    "        # retrieve text nodes\n",
    "        nodes = self.retriever.retrieve(query_str)\n",
    "        # create ImageNode items from text nodes\n",
    "        image_blocks = [\n",
    "            ImageBlock(path=n.metadata[\"image_path\"])\n",
    "            for n in nodes\n",
    "            if n.metadata.get(\"image_path\")\n",
    "        ]\n",
    "\n",
    "        # create context string from text nodes, dump into the prompt\n",
    "        context_str = \"\\n\\n\".join(\n",
    "            [r.get_content(metadata_mode=MetadataMode.LLM) for r in nodes]\n",
    "        )\n",
    "\n",
    "        formatted_msg = ChatMessage(\n",
    "            role=\"user\",\n",
    "            blocks=[\n",
    "                TextBlock(text=qa_prompt_block_text.format(context_str=context_str)),\n",
    "                image_prefix_block,\n",
    "                *image_blocks,\n",
    "                TextBlock(text=image_suffix.format(query_str=query_str)),\n",
    "            ],\n",
    "        )\n",
    "\n",
    "        # synthesize an answer from formatted text and images\n",
    "        llm_response = self.llm.chat([formatted_msg])\n",
    "\n",
    "        return Response(\n",
    "            response=str(llm_response.message.content),\n",
    "            source_nodes=nodes,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0890be59-fb12-4bb5-959b-b2d9600f7774",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = MultimodalQueryEngine(\n",
    "    retriever=index.as_retriever(similarity_top_k=3),\n",
    "    llm=Anthropic(model=\"claude-4-sonnet-20250514\"),\n",
    ")\n",
    "\n",
    "base_query_engine = base_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f94ef26-0df5-4468-a156-903d686f02ce",
   "metadata": {},
   "source": [
    "## Try out Queries\n",
    "\n",
    "Let's try out some questions against the slide deck in this multimodal RAG pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fd1aae3-1f8a-4797-a24a-17e563a7165e",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\n",
    "    \"which departments/teams use genAI the most and how are they using it?\"\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9cc48ee-481b-40b1-91b3-c69220e9dfb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_response = base_query_engine.query(\n",
    "    \"which departments/teams use genAI the most and how are they using it?\"\n",
    ")\n",
    "print(str(base_response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9462a82-960a-4c42-bbca-a1e71c2c1e5c",
   "metadata": {},
   "source": [
    "In this next question, the same sources are retrieved with and without contextual retrieval, and the answer is correct for both approaches. This is thanks for LlamaParse Premium's ability to comprehend graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8a0c8b1-3a3e-41c1-9916-01fdfb0dd8e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"what are relevant insights from the 'deep dive on infrastructure' section in terms of model preferences, cost, deployment environments?\"\n",
    "\n",
    "response = query_engine.query(query)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1638c6-ca29-462b-a21f-a2941968259c",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_response = base_query_engine.query(query)\n",
    "print(str(base_response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
